# coding:utf-8
# __author__ = yuan
# __time__ = 2020/2/25
# __file__ = quantization_1
# __desc__ =
import torch.quantization
import torch.nn as nn
import torch.jit
import os

"""
在准备量化模型时，有必要确保qconfig和用于量化计算的引擎与将执行模型的后端相匹配
backends: fbgemm for x86;qnnback for arm

在量化之前的操作：
  当前有必要在量化之前对模型定义进行一些修改. 这是因为当前的量化在逐个模块的基础上进行. 具体来说，对于所有量化技术，用户需要：
    将需要输出重新量化（因此具有其他参数）的所有操作从功能转换为模块形式.
    通过在子模块上分配.qconfig属性或通过指定qconfig_dict来指定需要量化模型的哪些部分
    对于量化激活的静态量化技术，用户还需要执行以下操作：
    指定对激活进行量化和反量化的位置. 这是使用QuantStub和DeQuantStub模块完成的.
    使用torch.nn.quantized.FloatFunctional将需要特殊处理以将量化量化的张量操作包装到模块中. 例如add和cat类的操作，它们需要特殊处理才能确定输出量化参数.
    保险丝模块：将操作/模块组合为一个模块，以获得更高的精度和性能. 这是使用torch.quantization.fuse_modules() API完成的，
    该API接收要融合的模块列表. 我们目前支持以下融合：[Conv，Relu]，[Conv，BatchNorm]，[Conv，BatchNorm，Relu]，[Linear，Relu]
    
torch.quantization
此模块实现您直接调用的功能，以将模型从FP32转换为量化形式. 例如， 
  prepare()用于后期训练量化，以为校准步骤准备模型， 
  convert()实际上将权重转换为int8，并用量化的对应项替换操作. 
  还有其他帮助程序功能，例如对模型的输入进行量化以及执行conv + relu等关键融合
"""


###  dynamic quantization
from torch import optim

#在大多数情况下，模型在FP32中进行训练，然后将模型转换为INT8。

quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.LSTM, nn.Linear}, dtype=torch.qint8
)
print(quantized_model)

### STATIC QUANTIZATION WITH EAGER MODE IN PYTORCH
'''
这个需要改网络的源代码
训练后静态量化; 及一般细节：
  准备模型：
      通过添加QuantStub和DeQuantStub模块，指定在何处显式量化激活和量化数量.
      Ensure that modules are not reused.
      将所有需要重新量化的操作转换为模块
  可以将诸如conv + relu或conv + batchnorm + relu之类的保险丝操作融合在一起，以提高模型的准确性和性能.
  指定'97量化方法的配置，例如选择对称或非对称量化以及MinMax或L2Norm校准技术.
  使用torch.quantization.prepare()插入将在校准期间观察激活张量的模块
  通过对校准数据集进行推断来校准模型
  最后，使用torch.quantization.convert（）方法转换模型本身. 这可以做几件事：量化权重，计算并存储每个激活张量要使用的比例和偏差值，并替换关键算子的量化实现
"""
 基于mobilenet的架构改进
      用nn.quantized.FloatFunctional代替加法
      在网络的开头和结尾处插入QuantStub和DeQuantStub .
      用ReLU替换ReLU6
    x = self.quant(x)
    x = self.features(x)
    x = x.mean([2, 3])
    x = self.classifier(x)
    x = self.dequant(x)
    return x
 myModel = load_model(saved_model_dir + float_model_file).to('cpu')
 myModel.eval()
 # Fuse Conv, bn and relu
 myModel.fuse_model()
"""
myModel = load_model(saved_model_dir + float_model_file).to('cpu')
myModel.eval()

# Fuse Conv, bn and relu
myModel.fuse_model() ## 这是自己写的Module里的函数

# Specify quantization configuration
# Start with simple min/max range estimation and per-tensor quantization of weights

myModel.qconfig = torch.quantization.default_qconfig
print(myModel.qconfig)

torch.quantization.prepare(myModel, inplace=True)
# Convert to quantized model

torch.quantization.convert(myModel, inplace=True)
'''
"""
首先通过网络馈送一批数据并计算不同激活的结果分布（具体而言，这是通过在记录此数据的不同点插入观察者模块来完成）. 
然后将这些分布用于确定在推理时如何具体量化不同的激活（一种简单的技术是将整个激活范围简单地划分为256个级别，但我们也支持更复杂的方法）. 
重要的是，此附加步骤使我们能够在操作之间传递量化值
"""
def post_train_staic_quantization():
    """
    Configuration for quantizing for x86 architectures. This configuration does the following:

    Quantizes weights on a per-channel basis
    Uses a histogram observer that collects a histogram of activations and then picks quantization parameters in an optimal manner.
    """
    per_channel_quantized_model = load_model(saved_model_dir + float_model_file)
    per_channel_quantized_model.eval()
    per_channel_quantized_model.fuse_model()
    per_channel_quantized_model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
    print(per_channel_quantized_model.qconfig)

    torch.quantization.prepare(per_channel_quantized_model, inplace=True)
    # 是否是根据evalute来进行校准的
    evaluate(per_channel_quantized_model,criterion, data_loader, num_calibration_batches)
    torch.quantization.convert(per_channel_quantized_model, inplace=True)
    top1, top5 = evaluate(per_channel_quantized_model, criterion, data_loader_test, neval_batches=num_eval_batches)
    print('Evaluation accuracy on %d images, %2.2f'%(num_eval_batches * eval_batch_size, top1.avg))
    torch.jit.save(torch.jit.script(per_channel_quantized_model), saved_model_dir + scripted_quantized_model_file)

######   Quantization-aware training  #####量化感知训练
"""
注意，整个计算是在浮点数中进行的。在量化感知训练的最后，
PyTorch提供了转换函数，将训练后的模型转换为精度较低的模型
by  FakeQuantize模块
计算将在FP32中进行，但将值取整并四舍五入以模拟INT8量化的效果. 步骤的顺序非常相似：
  步骤（1）和（2）相同.
  指定伪量化方法'97的配置，例如选择对称或非对称量化和MinMax或移动平均或L2Norm校准技术.
  使用torch.quantization.prepare_qat()插入将在训练过程中模拟量化的模块.
  训练或微调模型.
  与步骤（6）相同，用于训练后量化

在训练的正向和反向过程中，所有权重和激活都被"伪量化"：也就是说，
浮点值会四舍五入以模拟int8值，但所有计算仍将使用浮点数进行. 因此，在训练过程中进行所有权重调整，同时"意识到"模型将最终被量化的事实. 
因此，在量化之后，此方法通常比动态量化或训练后静态量化具有更高的准确性.

我们需要使用qconfig来指定要在权重和激活之后插入哪种伪量化，而不是指定观察者
"""

def qawaretrain():
    # 伪量化操作
    # We fuse modules as before

    qat_model = load_model(saved_model_dir + float_model_file)
    qat_model.fuse_model()

    optimizer = torch.optim.SGD(qat_model.parameters(), lr = 0.0001)
    qat_model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
    # Finally, prepare_qat performs the “fake quantization”, preparing the model for quantization-aware training
    # 最后， prepare_qat执行"伪量化"，为量化感知训练准备模型
    torch.quantization.prepare_qat(qat_model, inplace=True)
    print('Inverted Residual Block: After preparation for QAT, note fake-quantization modules \n',qat_model.features[1].conv)
    """
    对于量化感知训练，我们通过以下方式修改训练循环：
        在训练快要结束时切换批处理规范以使用运行均值和方差，以更好地匹配推理数字.
        我们还冻结了量化器参数（比例和零点），并对权重进行了微调.
    """
    for nepoch in range(8):
        train_one_epoch(qat_model, criterion, optimizer, data_loader, torch.device('cpu'), num_train_batches)
        if nepoch > 3:
            # Freeze quantizer parameters
            qat_model.apply(torch.quantization.disable_observer)
        if nepoch > 2:
            # Freeze batch norm mean and variance estimates
            qat_model.apply(torch.nn.intrinsic.qat.freeze_bn_stats)

        # Check the accuracy after each epoch
        quantized_model = torch.quantization.convert(qat_model.eval(), inplace=False)
        quantized_model.eval()
        top1, top5 = evaluate(quantized_model, criterion, data_loader_test, neval_batches=num_eval_batches)
        print(
          'Epoch %d :Evaluation accuracy on %d images, %2.2f' % (nepoch, num_eval_batches * eval_batch_size, top1.avg))


######### 基于迁移学习的量化模型 #######
import torchvision.models.quantization as models
# REW:目前只支持cpu上的量化训练 但是非量化部分可以转到gpu上
# You will need the number of filters in the `fc` for future use.
# Here the size of each output sample is set to 2.
# Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).
model_fe = models.resnet18(pretrained=True, progress=True, quantize=True)
num_ftrs = model_fe.fc.in_features
def create_combined_model(model_fe):
  # Step 1. Isolate the feature extractor.
  model_fe_features = nn.Sequential(
    model_fe.quant,  # Quantize the input
    model_fe.conv1,
    model_fe.bn1,
    model_fe.relu,
    model_fe.maxpool,
    model_fe.layer1,
    model_fe.layer2,
    model_fe.layer3,
    model_fe.layer4,
    model_fe.avgpool,
    model_fe.dequant,  # Dequantize the output
  )

  # Step 2. Create a new "head"
  new_head = nn.Sequential(
    nn.Dropout(p=0.5),
    nn.Linear(num_ftrs, 2),
  )

  # Step 3. Combine, and don't forget the quant stubs.
  new_model = nn.Sequential(
    model_fe_features,
    nn.Flatten(1),
    new_head,
  )
  return new_model
new_model = create_combined_model()
new_model.to('cpu')

# 在上部分中，我们使用了预训练的量化模型。在这一部分中，我们在对感兴趣的数据集进行微调后，创建了一个量化的特征提取器
#####    微调
# 预处理的特征提取器必须是可量化的。要确保它是可量化的，请执行以下步骤
    #1.融合(Conv,BN,Relu),(Conv,BN),(Conv,Relu)用torch.quantization.fuse_modules
    #2.用自定义头(分类器)连接特征提取器。这需要对特征提取器的输出进行去量化
    #3.在特征提取器的适当位置插入伪量化模块，以便在训练过程中模拟量化
# 对应上面步骤的实现
model = models.resnet18(pretrained=True, progress=True, quantize=False)
num_ftrs = model.fc.in_features
model.train()
model.fuse_model()
# Step 2
model_ft = create_combined_model(model)
model_ft[0].qconfig = torch.quantization.default_qat_qconfig # Use default QAT configuration
# Step 3
model_ft = torch.quantization.prepare_qat(model_ft, inplace=True)
for param in model_ft.parameters():
  param.requires_grad = True

model_ft.to('gpu')  # We can fine-tune on GPU if available

criterion = nn.CrossEntropyLoss()

# Note that we are training everything, so the learning rate is lower
# Notice the smaller learning rate
optimizer_ft = optim.SGD(model_ft.parameters(), lr=1e-3, momentum=0.9, weight_decay=0.1)

# Decay LR by a factor of 0.3 every several epochs
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=5, gamma=0.3)

model_ft_tuned = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
                             num_epochs=25, device=device)
# 转化
from torch.quantization import convert
model_ft_tuned.cpu()

model_quantized_and_trained = convert(model_ft_tuned, inplace=False)
# Step 1
model.train()
model.fuse_model()
# Step 2
model_ft = create_combined_model(model)
model_ft[0].qconfig = torch.quantization.default_qat_qconfig # Use default QAT configuration
# Step 3
model_ft = torch.quantization.prepare_qat(model_ft, inplace=True)
for param in model_ft.parameters():
  param.requires_grad = True

model_ft.to(device)  # We can fine-tune on GPU if available

criterion = nn.CrossEntropyLoss()

# Note that we are training everything, so the learning rate is lower
# Notice the smaller learning rate
optimizer_ft = optim.SGD(model_ft.parameters(), lr=1e-3, momentum=0.9, weight_decay=0.1)

# Decay LR by a factor of 0.3 every several epochs
exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=5, gamma=0.3)

model_ft_tuned = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
                             num_epochs=25, device=device)
# Step 5. Convert to quantized model

from torch.quantization import convert
model_ft_tuned.cpu()

model_quantized_and_trained = convert(model_ft_tuned, inplace=False)